package org.zjx.common.route;


import org.zjx.common.route.bean.Address;

import java.util.*;

/**
 * Dijkstra算法实现类（实现PathFinder接口）
 */
public class DijkstraPathFinder implements PathFinder {
    private static final Comparator<NodeTimePair> COMPARATOR =
            Comparator.comparingInt(ntp -> ntp.time);

    @Override
    public List<IRoute> findPath(List<IRoute> routes, Address start, Address end) {
        if (routes == null || routes.isEmpty()) return Collections.emptyList();
        if (start == null || end == null) return null;

        Map<Address, List<IRoute>> graph = buildGraph(routes);
        if (!graph.containsKey(start)) return null;

        Map<Address, Integer> distances = initializeDistances(routes);
        Map<Address, PathNode> pathMap = new HashMap<>();
        PriorityQueue<NodeTimePair> pq = new PriorityQueue<>(COMPARATOR);

        executeDijkstra(graph, start, distances, pathMap, pq);
        return reconstructPath(end, pathMap, start);
    }

    private Map<Address, List<IRoute>> buildGraph(List<IRoute> routes) {
        Map<Address, List<IRoute>> graph = new HashMap<>();
        for (IRoute route : routes) {
            if (route == null) continue;
            graph.computeIfAbsent(route.getStart(), k -> new ArrayList<>()).add(route);
        }
        return graph;
    }

    private Map<Address, Integer> initializeDistances(List<IRoute> routes) {
        Map<Address, Integer> distances = new HashMap<>();
        Set<Address> allNodes = new HashSet<>();
        for (IRoute route : routes) {
            allNodes.add(route.getStart());
            allNodes.add(route.getEnd());
        }
        for (Address node : allNodes) {
            distances.put(node, Integer.MAX_VALUE);
        }
        return distances;
    }

    private void executeDijkstra(Map<Address, List<IRoute>> graph, Address start,
                                 Map<Address, Integer> distances, Map<Address, PathNode> pathMap,
                                 PriorityQueue<NodeTimePair> pq) {
        distances.put(start, 0);
        pq.offer(new NodeTimePair(start, 0));

        while (!pq.isEmpty()) {
            NodeTimePair current = pq.poll();
            Address currentNode = current.node;
            int currentDist = current.time;

            if (currentDist > distances.getOrDefault(currentNode, Integer.MAX_VALUE)) continue;

            if (!graph.containsKey(currentNode)) continue;

            for (IRoute route : graph.get(currentNode)) {
                Address nextNode = route.getEnd();
                int newDist = currentDist + route.getTime();

                if (newDist < distances.getOrDefault(nextNode, Integer.MAX_VALUE)) {
                    distances.put(nextNode, newDist);
                    pathMap.put(nextNode, new PathNode(currentNode, route));
                    pq.offer(new NodeTimePair(nextNode, newDist));
                }
            }
        }
    }

    private List<IRoute> reconstructPath(Address end, Map<Address, PathNode> pathMap, Address start) {
        List<IRoute> path = new ArrayList<>();
        Address current = end;

        while (current != null) {
            PathNode nodeInfo = pathMap.get(current);
            if (nodeInfo == null) break;
            path.add(nodeInfo.route);
            current = nodeInfo.prevNode;
        }

        if (path.isEmpty() || !path.get(path.size() - 1).getStart().equals(start)) {
            return null;
        }

        Collections.reverse(path);
        return path;
    }

    private static class NodeTimePair {
        Address node;
        int time;

        NodeTimePair(Address node, int time) {
            this.node = node;
            this.time = time;
        }
    }

    private static class PathNode {
        Address prevNode;
        IRoute route;

        PathNode(Address prevNode, IRoute route) {
            this.prevNode = prevNode;
            this.route = route;
        }
    }
}